#!/usr/bin/python
# Copyright (C) 2013, 2014 Christian Sengstock
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Frequent itemset mining with FP-Growth

Usage:

>> from fpgrowth import fim
>> tdb = [('a','b','c'), ('a','b'), ('b','c','d','e')]
>> mincount = 2
>> res = fim(tdb, mincount)
>> print(res)
   [(('b',), 3), (('a',), 2), (('c',), 2), (('b', 'a'), 2), (('b', 'c'), 2)]

About:

FP-Growth is a divide and conquer algorithm for frequent itemset mining
proposed in:

  Han J., Pei H., Yin Y. Mining Frequent Patterns
  without Candidate Generation. In: Proc. Conf. on
  the Management of Data (SIGMOD '00).

Code: 

Pure Python implemenation with some performance optimizations."""


__author__ = "csengstock@gmail.com (Christian Sengstock)"
__version__ = 1.1

import sys
from collections import defaultdict
import itertools

VERB_OUT = sys.stderr
VERB = False

def _vout(s):
    if VERB: VERB_OUT.write(s)


def fim(tdb, mincount, spp_opt=True):
    """
    Frequent itemset mining on a transaction database (tdb).

    Args:
        tdb:       Iterable of iterables of items. Items need
                    to be hashable
        mincount:  Minimum count threshold (inclusive)
        ssp_opt:   Use 'single prefix path' optimization (default: True)
    """

    _vout("tdb stats...\n")

    item_counts, n_tdb, n_items = \
         _tdb_stats(tdb, mincount)
       
    _vout("# of transactions: %d\n" % n_tdb)
    _vout("# of items:        %d\n" % n_items)
    _vout("# of freq items:   %d\n" % len(item_counts))
    _vout("mincount:          %d\n" % mincount)
    _vout("minsup:            %.4f\n" % (mincount/float(n_tdb)))
    _vout("fptree...\n")

    fptree = FPTree(item_counts)
    for t in tdb:
        fptree.add(t)

    _vout("fpgrowth...\n")

    return fptree.fpgrowth(mincount, spp_opt)


def _tdb_stats(tdb, mincount):
    item_counts = defaultdict(int)
    n_tdb = 0
    for t in tdb:
        n_tdb += 1
        s = set(t)
        for item in s:
            item_counts[item] += 1
    n_items = len(item_counts)
    for item in item_counts.keys():
        if item_counts[item] < mincount:
            del item_counts[item]
    return item_counts, n_tdb, n_items


def powerset(l):
    return itertools.chain.from_iterable(
            itertools.combinations(l, r) for r in xrange(1,len(l)+1))

def powersetminer(tdb, mincount):
    """For testing only."""

    items = set([])
    for t in tdb:
        for item in t:
            items.add(item)

    itemsets = powerset(list(items))
    counts = {}
    for itemset in itemsets:
        s = set(itemset)
        k = list(itemset)
        k.sort()
        k = "".join(str(i) for i in k)
        for t in tdb:
            if len(s - set(t)) == 0:
                counts.setdefault(k,[tuple(itemset),0])
                counts[k][1] += 1

    res = []
    for k, itemset_count in counts.items():
        if itemset_count[1] >= mincount:
            res.append( itemset_count )
    return res


class FPNode(object):
    __slots__ = ["item", "count", "parent", "children"]

    def __init__(self, item, count, parent, children):
        self.item = item
        self.count = count
        self.parent = parent
        self.children = children


class FPTree(object):
    __slots__ = ["alpha", "root", "head_nodes", "head_counts", "head_items"]

    def __init__(self, item_counts, alpha=()):
        """Constructs a new FP-Tree for the items
        given in the item_counts dictionary.

        Args:
            item_counts:   (item => count) dictionary
            alpha:         FP-Tree prefix
        """

        self.head_nodes = defaultdict(list)
        self.head_counts = item_counts
        self.head_items = set(item_counts.keys())
         
        self.alpha = alpha
        self.root = FPNode(None, None, None, {})

    def add(self, itemset, count=1):

        # filter valid items and bring
        #  itemset into canonical order.

        itemset = list(self.head_items & set(itemset))
        if len(itemset) == 0:
            return

        itemset.sort(key=lambda x: (self.head_counts[x], x), reverse=True)

        # add nodes to the tree

        parent = self.root
        for item in itemset:
            if item in parent.children:
                child = parent.children[item]
                child.count += count
                parent = child
            else:
                child = FPNode(item, count, parent, {})
                self.head_nodes[child.item].append(child)
                parent.children[item] = child
                parent = child

    def _cond_pattern_base(self, item, mincount=1):
        item_counts = defaultdict(int)
        cond_pattern_base = []

        # collect paths from items' linknodes to root node.

        for linknode in self.head_nodes[item]:
            itemset = []
            pathnode = linknode.parent
            while pathnode.parent != None:
                itemset.append(pathnode.item)
                item_counts[pathnode.item] += linknode.count
                pathnode = pathnode.parent
            if len(itemset) > 0:
                cond_pattern_base.append( (itemset, linknode.count) )

        # remove in-frequent items from count dict.

        for item in item_counts.keys():
            if item_counts[item] < mincount:
                del item_counts[item]

        return cond_pattern_base, item_counts

    def _cond_pattern_tree(self, item, mincount=1):
        cond_pattern_base, item_counts = \
            self._cond_pattern_base(item, mincount)

        if not cond_pattern_base:
            return None

        fptree = FPTree(item_counts, self.alpha + (item,))

        for itemset, count in cond_pattern_base:
            fptree.add(itemset, count)

        return fptree

    def fpgrowth(self, mincount, pfp_opt=True):
        if pfp_opt:
            return self._fpgrowth_opo(mincount)
        else:         
            result = []
            self._fpgrowth(mincount, result)
            return result
   
    def _fpgrowth(self, mincount, result):

        # Algorithm:
        #
        # 1) Generate frequent itemsets by concatenating the 
        #    conditional itemset (alpha) with the frequent
        #    items in the header table.
        #
        # 2) Generate the conditional pattern trees
        #    for the frequent items in the header table
        #    and call fpgrowth (recursion).        
     
        for item, count in self.head_counts.iteritems():
            # 1) ...
            assert count >= mincount
            result.append( (self.alpha + (item,), count) )

            # 2) ...
            cond_pattern_tree = self._cond_pattern_tree(
                        item, mincount)

            # If tree only consists of a root node,
            # skip recursive call
            if cond_pattern_tree:
                subresult = cond_pattern_tree._fpgrowth(mincount, result)

    def _fpgrowth_opo(self, mincount):
        """
        fpgrowth with 'single prefix path' optimization.
        """

        prefix_path, subtree = self._extract_prefix_path()
        prefix_result = self._prefix_path_result(prefix_path, mincount)
        tree_result = self._subtree_result(subtree, mincount)
     
        # all unions between prefix path and tree result
        comb_result = []
        if len(tree_result) != 0 and len(prefix_result) != 0:
            for prefix_itemset in prefix_result:
                for tree_itemset in tree_result:
                    union_itemset = tuple( set(prefix_itemset[0]) | \
                                           set(tree_itemset[0]) )
                    comb_result.append( (union_itemset, min(tree_itemset[1],\
                                                            prefix_itemset[1])) )

        prefix_result.extend(tree_result)
        prefix_result.extend(comb_result)
        return prefix_result

    def _extract_prefix_path(self):

        # Collect prefix path itemset
        # starting from root.

        prefix_path = []
        node = self.root
        while len(node.children) == 1:
            node = node.children.values()[0]
            prefix_path.append( (node.item, node.count) )

        if len(prefix_path) == 0:
            return prefix_path, self

        subtree = None
        if len(node.children) > 1:
            subtree = FPTree({}, self.alpha)

            subtree.root.children = node.children
            for item, child in subtree.root.children.iteritems():
                child.parent = subtree.root

            # Update the subtree header table.
            # Items in the prefix path are not 
            # in the subtree. Insert the valid head
            # node paths from the old trees' header table.
            # item_set need not to be updated since
            # no more itemsets are added to this tree.

            prefix_items = set(zip(*prefix_path)[0])

            for item, nodes in self.header.items.iteritems():
                if item not in prefix_items:
                    subtree.header.items[item] = nodes
                    subtree.header.item_counts[item] = self.header.item_counts[item]

        return prefix_path, subtree
  
    def _prefix_path_result(self, prefix_path, mincount):
        prefix_result = []
        if len(prefix_path) > 0:
            prefix_powerset = powerset(prefix_path)
            for prefix_comb in prefix_powerset:
                min_count = prefix_comb[0][1]
                itemset = []
                for item, count in prefix_comb:
                    if count < min_count: min_count = count
                    itemset.append(item)
                if min_count >= mincount:
                    prefix_result.append( (self.alpha + tuple(itemset), count) )
        return prefix_result

    def _subtree_result(self, subtree, mincount):
        subtree_result = []

        # see _fpgrowth() for comments. Code has same structure.
        if subtree != None:
            for item, count in subtree.head_counts.iteritems():
                assert count >= mincount
                subtree_result.append( (self.alpha + (item,), count) )
                cond_pattern_tree = subtree._cond_pattern_tree(item, mincount)
                if cond_pattern_tree:
                    subresult = cond_pattern_tree._fpgrowth_opo(mincount)
                    subtree_result.extend(subresult)
        return subtree_result



# /////////////////////////////////////////////////////////////////


def __perf_test():
    global tdb
    tdb = [list(s) for s in ["abcde", "abcfg", "abchi", "abc", "ab","a"]]
    from profile import run
    run("fim(tdb, 2, False)")
    run("fim(tdb, 2, True)")

def __test():   
    tdb = [list(s) for s in ["abcde", "abcfg", "abchi", "abc", "ab","a"]]
    res = fim(tdb, 2)
    res.sort(key=lambda x: (x[1], len(x[0])), reverse=True)
    for r in res:
        print r
    print ""
    res = _powersetminer(tdb, 2)
    res.sort(key=lambda x: (x[1], len(x[0])), reverse=True)
    for r in res:
        print r
  
if __name__ == "__main__":
    #__perf_test()
    __test()
